from tensorflow.keras.layers import Conv2D, MaxPooling2D, Activation, Flatten, Dense
from tensorflow.keras.models import Sequential
from tensorflow.keras import backend


class Lenet:
    @staticmethod
    def build(*, width: int, height: int, depth: int,
              classes: int) -> Sequential:
        """
        模型构建

        :param width: 图片宽
        :param height: 图片高
        :param depth: 图片通道数
        :param classes: 类别数
        :return: 模型对象
        """
        # 初始化模型
        model = Sequential()
        input_shape = (height, width, depth) if backend.image_data_format() != "channels_first" \
            else (depth, height, width)
        model.add(Conv2D(20, (5, 5), padding="same",
                         input_shape=input_shape))
        model.add(Activation("relu"))
        model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

        # 再次添加
        model.add(Conv2D(50, (5, 5), padding="same"))
        model.add(Activation("relu"))
        model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

        # 展平输入
        model.add(Flatten())
        model.add(Dense(500))
        model.add(Activation("relu"))
        model.add(Dense(classes))
        model.add(Activation("softmax"))

        return model


if __name__ == '__main__':
    Lenet.build(width=28, height=28, depth=3, classes=10)
    # import tensorflow as tf
    # input_shape = (4, 28, 28, 3)
    # x = tf.random.normal(input_shape)
    # y = tf.keras.layers.Conv2D(20, (5, 5), input_shape=input_shape[1:], padding="same")(x)
    # print(y.shape)
